"""
解析PDF电子发票，提取信息
"""

import os
import fitz
import re
import cv2
import pandas as pd
import numpy as np
from PIL import Image, ImageDraw
from collections import defaultdict as def_dict
from .logger import get_logger

__author__ = 'yczha'
__email__ = 'yooongchun@foxmail.com'


logger = get_logger('INFO', save_path=__file__)


class InvoiceParser:
    """提取给定路径的电子发票pdf文件内容"""

    def __init__(self, path):
        """
        path: pdf文件路径
        """
        self.file = path
        self.file_name = os.path.basename(self.file)
        self._check_file()

    def _check_file(self):
        """检查文件是否合法"""
        if not isinstance(self.file, str) or not os.path.isfile(self.file) or not self.file.endswith(('.pdf', '.PDF')):
            raise TypeError('不是有效的pdf文件!!!')

    @staticmethod
    def _merge_block_text_into_line(block, delta=6):
        """功能：将block内容块中的内容分行
        block:内容块
        delta：同一行的y可偏移范围
        """
        words_in_line = def_dict(list)
        key = None
        ordered_words = sorted(block, key=lambda x: x[1])
        for word in ordered_words:
            y_val = int(round(word[1]))
            if key is not None and y_val in key:
                words_in_line[key].append(word)
            else:
                key = tuple([(i + y_val) for i in range(delta)])
                words_in_line[key].append(word)
        for key, words in words_in_line.items():
            str_ = ''.join(word[4] for word in sorted(words, key=lambda x: x[0]))
            str_ = str_.replace(' ', '').replace('\n', '')
            words_in_line[key] = str_

        text = ''
        words_in_line_new = {}
        for idx, (_, str_) in enumerate(sorted(words_in_line.items(), key=lambda x: x[0][0])):
            words_in_line_new[idx] = str_
            text += str_
        words_in_line_new['text'] = text
        return words_in_line_new

    @staticmethod
    def _split_blocks(page, line_delta=5):
        """功能：将页面内容根据特定单词位置分割，分割为9个内容块返回
        page: fitz加载的页面对象
        line_delta: 行外延距离
        """
        _, _, width, height = page.rect
        r1 = page.searchFor('密')[0]
        r2 = page.searchFor('税率')[0]
        r3 = page.searchFor('价税合计')[0]
        r4 = page.searchFor('款')[0]
        r5 = page.searchFor('纳税人识别号')[0]
        r6 = None
        for w1 in page.searchFor('合'):
            for w2 in page.searchFor('计'):
                if abs(w1[1] - w2[1]) < 0.1 < abs(w1[1] - r3[1]):
                    r6 = w1
                    break
        r7, r8 = sorted(page.searchFor('开户行及账号'), key=lambda x: x[1])[:2]

        blocks = {'rect1': fitz.Rect(0, 0, width, r1[1] - line_delta),  # 发票头部内容，包含矩形框以上所有内容
                  'rect2': fitz.Rect(r7[0] - line_delta, r1[1] - line_delta, r1[0] - line_delta, r7[3] + line_delta),
                  # 购买方
                  'rect3': fitz.Rect(r1[2], r1[1] - line_delta, width, r7[3] + line_delta),  # 密码区
                  'rect4': fitz.Rect(0, r2[1] - line_delta, width, r6[1] - line_delta),  # 详情栏，不含合计行
                  'rect5': fitz.Rect(0, r6[1] - line_delta, width, r6[3] + line_delta),  # 合计行
                  'rect6': fitz.Rect(0, r3[1] - line_delta, width, r3[3] + line_delta),  # 价税合计栏
                  'rect7': fitz.Rect(r5[0] - line_delta, r3[3] + line_delta, r1[0] - line_delta, r8[3] + line_delta),
                  # 销售方
                  'rect8': fitz.Rect(r1[2], r3[3] + line_delta, width, r8[3] + line_delta),  # 备注
                  'rect9': fitz.Rect(0, r4[1] - line_delta, width, height)}  # 尾栏

        return blocks

    def _put_words_into_blocks(self, page):
        """功能：将文字放入其对应的内容块(block)中
        page: fitz 加载的页面对象
        """
        page_blocks = self._split_blocks(page)
        text_blocks = def_dict(list)
        words = page.getText('words')
        for word in words:
            text_rect = fitz.Rect(word[:4])
            for blo_name, rect in page_blocks.items():
                if rect.contains(text_rect):
                    text_blocks[blo_name].append(word)
        for block in text_blocks:
            logger.debug(f'{block}:\t{[word[4] for word in text_blocks[block]]}')
        return page_blocks, text_blocks

    @staticmethod
    def _find_mid_sep_by_words_gap(block, rect, sep_num):
        """
        功能：查找行空隙，按照空隙大小分割为指定n行
        参数：
            block:内容块
            rect:内容块的矩形信息
            sep_num:查找的线段数目
        返回：分割之后的x坐标序列
        """
        img = Image.new('RGB', [int(rect.x1), int(rect.y1)])
        draw = ImageDraw.Draw(img)
        for word in block:
            draw.rectangle(fitz.Rect(word[:4]), outline='blue', width=2)
        row_sum = np.array(img.convert('L')).sum(axis=0)
        sep_list = []
        for curr in range(1, row_sum.shape[0]):
            prev = curr - 1
            if row_sum[prev] > 0 and row_sum[curr] == 0:
                sep_list.append([curr, None])
            elif sep_list and row_sum[prev] == 0 and row_sum[curr] > 0:
                sep_list[-1][1] = prev
        sep_list = [[v1, v2] for v1, v2 in sep_list if v1 and v2]
        sep_list = sorted(sep_list, key=lambda x: x[1] - x[0], reverse=True)[:sep_num]
        sep_mid = sorted((v1 + v2) // 2 for v1, v2 in sep_list)
        return [0] + sep_mid + [1000]

    @staticmethod
    def _find_mid_sep_by_cv_hough_line(page, rect, sep_num):
        """
        功能：查找分割行
        参数：
            page:fitz.page对象
            rect:内容块的矩形信息
            sep_num：要查找的线段数
        返回：分割之后的x坐标序列
        """
        pix = page.getPixmap(clip=rect)
        img = np.array(tuple(pix.samples), dtype=np.uint8).reshape([pix.height, pix.width, -1])
        gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray_img, 50, 180, apertureSize=3)
        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, 60, minLineLength=50, maxLineGap=5)
        vlines = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            if abs(x2 - x1) < 1:  # 挑选竖线
                vlines.append(line[0])
        vlines_sort = sorted(vlines, key=lambda x: x[3] - x[1], reverse=True)  # 按长度逆排序
        vlines_sort = sorted(vlines_sort, key=lambda x: x[0])  # 按横坐标排序
        vlines_drop = [vlines_sort[0]]
        line_gap = 5
        for line in vlines_sort[1:]:
            if abs(line[0] - vlines_drop[-1][0]) >= line_gap:
                vlines_drop.append(line)
        sep_list = [line[0] for line in vlines_drop[:sep_num]]
        return sep_list

    def _split_column_blocks(self, page, block, rect, sep_row_num=8):
        """
        功能：将内容块按列划分为多个子块
        参数：
            block：内容块
            rect:内容块的坐标
            sep_row_num：需要划分的子块数目
        """
        sep_mid = self._find_mid_sep_by_cv_hough_line(page, rect, sep_row_num + 1)
        # sep_mid = self._find_mid_sep_by_words_gap(block, rect, sep_row_num-1)
        blocks_col = def_dict(list)
        for curr in range(1, len(sep_mid)):
            prev = curr - 1
            for word in block:
                x = (word[0] + word[2]) / 2
                if sep_mid[prev] <= x <= sep_mid[curr]:
                    blocks_col[sep_mid[curr]].append(word)
        return blocks_col

    def _search_block_head(self, block):
        """
        功能：查找头部内容块的内容，包含：
            1.发票标题
            2.发票号码
            3.发票代码
            4.开票日期
            5.机器编号
            6.校验码
        """
        block = self._merge_block_text_into_line(block)
        text = block['text']
        # 省份名称列表，为了查表发票标题使用
        province_names = ['河北', '山西', '辽宁', '吉林', '黑龙江', '江苏', '浙江', '安徽', '福建', '江西', '山东',
                          '河南', '湖北', '湖南', '广东', '海南', '四川', '贵州', '云南', '陕西', '甘肃', '青海',
                          '台湾', '北京', '天津', '上海', '重庆', '内蒙古', '广西', '西藏', '宁夏', '新疆']
        prov_names = [prov for prov in province_names if prov in text]
        try:
            start, end = re.search(r'%s.*普通发票|%s.*专用发票' % (prov_names[0], prov_names[0]), text).span()
            inv_title = text[start:end]
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_title = ''
        try:
            start, end = re.search(r'发票代码[:：\s]*\d+', text).span()
            inv_code = text[start:end]
            inv_code = re.sub(r'发票代码[:：\s]+', '', inv_code)
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_code = ''
        try:
            start, end = re.search(r'发票号码[:：\s]*\d+', text).span()
            inv_num = text[start:end]
            inv_num = re.sub(r'发票号码[:：\s]+', '', inv_num)
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_num = ''
        try:
            start, end = re.search(r'开票日期[:：\s]*\d{4}年\d{1,2}月\d{1,2}日', text).span()
            inv_date = text[start:end]
            inv_date = re.sub(r'开票日期[:：\s]+', '', inv_date)
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_date = ''
        try:
            start, end = re.search(r'校验码[:：\s]*\d+', text).span()
            inv_check_code = text[start:end]
            inv_check_code = re.sub(r'校验码[:：\s]+', '', inv_check_code)
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_check_code = ''
        try:
            start, end = re.search(r'机器编号[:：\s]*\d+', text).span()
            inv_machine_code = text[start:end]
            inv_machine_code = re.sub(r'机器编号[:：\s]+', '', inv_machine_code)
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_machine_code = ''
        return {'inv_title': inv_title, 'inv_code': inv_code, 'inv_num': inv_num, 'inv_date': inv_date,
                'inv_check_code': inv_check_code, 'inv_machine_code': inv_machine_code}

    def _search_block_buyer(self, block):
        """
        功能：查找购买方信息，包含：
            1.名称
            2.纳税人识别号
            3.地址、电话
            4.开户行及帐号
        """
        block = self._merge_block_text_into_line(block)
        text = block['text']
        text = re.sub(r'[:：\s\n、]+', '', text)
        vli = re.split(r'名称|纳税人识别号|开户行及账号|地址电话', text) + ['', '', '', '', '']
        try:
            inv_buyer_name, inv_buyer_id, inv_buyer_addr, inv_buyer_account = vli[1:5]
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_buyer_name, inv_buyer_id, inv_buyer_addr, inv_buyer_account = '', '', '', ''
        return {'inv_buyer_name': inv_buyer_name, 'inv_buyer_id': inv_buyer_id,
                'inv_buyer_address': inv_buyer_addr, 'inv_buyer_account': inv_buyer_account}

    def _search_block_password(self, block):
        """
        功能：查找密码区
        """
        block = self._merge_block_text_into_line(block)
        return block['text']

    def _search_block_details(self, page, block, rect):
        """
        功能：查找详细信息内容块
        参数：
            block：内容块
            rect：内容块对应的矩形框
        """
        col_blocks = self._split_column_blocks(page, block, rect, sep_row_num=8)
        col_blocks_sorted = sorted(col_blocks.items(), key=lambda x: x[0])
        col_block_words = []
        for _, val in col_blocks_sorted:
            col_block_lines = self._merge_block_text_into_line(val)
            col_block_lines.pop('text')
            col_block_lines_items = sorted(col_block_lines.items(), key=lambda x: x[0])
            col_block_lines_words = [word for _, word in col_block_lines_items]
            col_block_words.append(col_block_lines_words)
        try:
            inv_proj_words = col_block_words[0]
            inv_detail_first_name = inv_proj_words[0]
            if len(inv_proj_words) > 1:
                inv_detail_first_val = '\n'.join(str(val) for val in inv_proj_words[1:])
            else:
                inv_detail_first_val = ''
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_detail_first_name = 'detail1_未找到'
            inv_detail_first_val = ''
        try:
            inv_sec_words = col_block_words[1]
            inv_detail_second_name = inv_sec_words[0]
            if len(inv_sec_words) > 1:
                inv_detail_second_val = '\n'.join(str(val) for val in inv_sec_words[1:])
            else:
                inv_detail_second_val = ''
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_detail_second_name = 'detail2_未找到'
            inv_detail_second_val = ''
        try:
            inv_third_words = col_block_words[2]
            inv_detail_third_name = inv_third_words[0]
            if len(inv_third_words) > 1:
                inv_detail_third_val = '\n'.join(str(val) for val in inv_third_words[1:])
            else:
                inv_detail_third_val = ''
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_detail_third_name = 'detail3_未找到'
            inv_detail_third_val = ''
        try:
            inv_four_words = col_block_words[3]
            inv_detail_fourth_name = inv_four_words[0]
            if len(inv_four_words) > 1:
                inv_detail_fourth_val = '\n'.join(str(val) for val in inv_four_words[1:])
            else:
                inv_detail_fourth_val = ''
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_detail_fourth_name = 'detail4_未找到'
            inv_detail_fourth_val = ''
        try:
            inv_fifth_words = col_block_words[4]
            inv_detail_fifth_name = inv_fifth_words[0]
            if len(inv_fifth_words) > 1:
                inv_detail_fifth_val = '\n'.join(str(val) for val in inv_fifth_words[1:])
            else:
                inv_detail_fifth_val = ''
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_detail_fifth_name = 'detail5_未找到'
            inv_detail_fifth_val = ''
        try:
            inv_sixth_words = col_block_words[5]
            if len(inv_sixth_words) > 1:
                inv_detail_sixth = '\n'.join(str(val) for val in inv_sixth_words[1:])
            else:
                inv_detail_sixth = ''
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_detail_sixth = ''
        try:
            inv_seventh_words = col_block_words[6]
            if len(inv_seventh_words) > 1:
                inv_detail_seventh = '\n'.join(str(val) for val in inv_seventh_words[1:])
            else:
                inv_detail_seventh = ''
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_detail_seventh = ''
        try:
            inv_eighth_words = col_block_words[7]
            if len(inv_eighth_words) > 1:
                inv_detail_eighth = '\n'.join(str(val) for val in inv_eighth_words[1:])
            else:
                inv_detail_eighth = ''
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_detail_eighth = ''

        return {'inv_detail_first_name': inv_detail_first_name,
                'inv_detail_first_val': inv_detail_first_val,
                'inv_detail_second_name': inv_detail_second_name,
                'inv_detail_second_val': inv_detail_second_val,
                'inv_detail_third_name': inv_detail_third_name,
                'inv_detail_third_val': inv_detail_third_val,
                'inv_detail_fourth_name': inv_detail_fourth_name,
                'inv_detail_fourth_val': inv_detail_fourth_val,
                'inv_detail_fifth_name': inv_detail_fifth_name,
                'inv_detail_fifth_val': inv_detail_fifth_val,
                'inv_detail_money': inv_detail_sixth,
                'inv_detail_tax_rate': inv_detail_seventh,
                'inv_detail_tax_num': inv_detail_eighth}

    def _search_block_summation(self, block):
        """
        功能：搜索合计栏
        """
        # TODO: 优化搜索方式
        block = self._merge_block_text_into_line(block)
        text = block['text']
        vli = re.findall(r'[¥￥][\d.]+', text) + ['', '']
        try:
            inv_sum_num = vli[0]
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_sum_num = ''
        try:
            inv_sum_rate = vli[1]
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_sum_rate = ''
        return {'inv_summation_money': inv_sum_num, 'inv_summation_tax_num': inv_sum_rate}

    def _search_block_summation_with_tax(self, block):
        """
        功能：搜索价税合计栏
        """
        block = self._merge_block_text_into_line(block)
        text = block['text']
        try:
            start, end = re.search(r'价税合计[(（]+大写[)）].*[（(]+小写', text).span()
            inv_sum_up = re.sub(r'价税合计|大写|小写|[(（)）]', '', text[start:end])
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_sum_up = ''
        try:
            start, end = re.search(r'[（(]+小写[)）].*', text).span()
            inv_sum_low = re.sub(r'小写|[(（)）]', '', text[start:end])
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_sum_low = ''
        return {'inv_tax_up': inv_sum_up, 'inv_tax_low': inv_sum_low}

    def _search_block_seller(self, block):
        """
        功能：查找销售方信息，包含：
            1.名称
            2.纳税人识别号
            3.地址、电话
            4.开户行及帐号
        """
        block = self._merge_block_text_into_line(block)
        text = block['text']
        text = re.sub(r'[:：\s\n、]+', '', text)
        vli = re.split(r'名称|纳税人识别号|开户行及账号|地址电话', text) + ['', '', '', '', '']
        try:
            inv_seller_name, inv_seller_id, inv_seller_addr, inv_seller_account = vli[1:5]
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_seller_name, inv_seller_id, inv_seller_addr, inv_seller_account = '', '', '', ''
        return {'inv_seller_name': inv_seller_name, 'inv_seller_id': inv_seller_id,
                'inv_seller_address': inv_seller_addr, 'inv_seller_account': inv_seller_account}

    def _search_block_note(self, block):
        """
        功能：搜索备注栏
        """
        block = self._merge_block_text_into_line(block)
        return block['text']

    def _search_block_tail(self, block):
        """
        功能：搜索尾部信息
        """
        block = self._merge_block_text_into_line(block)
        text = block['text']
        text = re.sub(r'[:：\s]+', '', text)
        vli = re.split(r'收款人|复核|开票人|销售方', text) + ['', '', '', '']
        try:
            inv_charge, inv_checker, inv_report, inv_seller = vli[1:5]
        except Exception as e:
            logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:{e}')
            inv_charge, inv_checker, inv_report, inv_seller = '', '', '', ''
        return {'inv_charge': inv_charge, 'inv_checker': inv_checker,
                'inv_poster': inv_report, 'inv_seller': inv_seller}

    def do_parse(self):
        columns = {'inv_title': '发票标题', 'inv_code': '发票代码', 'inv_num': '发票号码', 'inv_date': '开票日期',
                   'inv_check_code': '校验码', 'inv_machine_code': '机器编号',
                   'inv_buyer_name': '购买方_名称', 'inv_buyer_id': '购买方_纳税人识别号', 'inv_buyer_address': '购买方_地址电话',
                   'inv_buyer_account': '购买方_开户行及账号', 'inv_password': '密码区',
                   'inv_detail_first_name': '项目名称', 'inv_detail_first_val': '',
                   'inv_detail_second_name': '车牌号', 'inv_detail_second_val': '',
                   'inv_detail_third_name': '类型', 'inv_detail_third_val': '',
                   'inv_detail_fourth_name': '通行日期起', 'inv_detail_fourth_val': '',
                   'inv_detail_fifth_name': '通行日期止', 'inv_detail_fifth_val': '',
                   'inv_detail_money': '金额', 'inv_detail_tax_rate': '税率', 'inv_detail_tax_num': '税额',
                   'inv_summation_money': '合计_金额', 'inv_summation_tax_num': '合计_税额',
                   'inv_tax_up': '价税合计(大写)', 'inv_tax_low': '价税合计(小写)', 'inv_note': '备注',
                   'inv_seller_name': '销售方_名称', 'inv_seller_id': '销售方_纳税人识别号',
                   'inv_seller_address': '销售方_地址电话', 'inv_seller_account': '销售方_开户行及账号',
                   'inv_charge': '收款人', 'inv_checker': '复核', 'inv_poster': '开票人', 'inv_seller': '销售方'}
        df = pd.DataFrame(columns=list(columns.keys()))
        doc = fitz.open(self.file)
        for idx, page in enumerate(doc):
            if not page.searchFor('普通发票') and not page.searchFor('专用发票'):
                logger.error(f'[pid:{os.getpid()} file:{self.file_name}] msg:无效的发票 page:{idx + 1}')
                continue
            page_idx = f'page_{idx + 1}'
            rects, blocks = self._put_words_into_blocks(page)
            # block1
            head = self._search_block_head(blocks['rect1'])
            for name, val in head.items():
                df.loc[page_idx, name] = val
            # block2
            buyer = self._search_block_buyer(blocks['rect2'])
            for name, val in buyer.items():
                df.loc[page_idx, name] = val
            # block3
            password = self._search_block_password(blocks['rect3'])
            df.loc[page_idx, 'inv_password'] = password
            # block4
            inv_details = self._search_block_details(page, blocks['rect4'], rects['rect4'])
            for key, val in inv_details.items():
                df.loc[page_idx, key] = val
            # block5
            inv_summation = self._search_block_summation(blocks['rect5'])
            for key, val in inv_summation.items():
                df.loc[page_idx, key] = val
            # block6
            inv_summation_tax = self._search_block_summation_with_tax(blocks['rect6'])
            for key, val in inv_summation_tax.items():
                df.loc[page_idx, key] = val
            # block7
            seller = self._search_block_seller(blocks['rect7'])
            for name, val in seller.items():
                df.loc[page_idx, name] = val
            # block8
            inv_note = self._search_block_note(blocks['rect8'])
            df.loc[page_idx, 'inv_note'] = inv_note
            # block9
            inv_tail = self._search_block_tail(blocks['rect9'])
            for name, val in inv_tail.items():
                df.loc[page_idx, name] = val

        return df


if __name__ == '__main__':
    parser = InvoiceParser('example/test.pdf')
    result = parser.do_parse()
    print(result.T)
